{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "219f48a4",
   "metadata": {},
   "source": [
    "# TREx API Examples\n",
    "\n",
    "This notebook provides a set of TREx API examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c16a4d6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import trex\n",
    "from trex.notebook import *\n",
    "\n",
    "engine_name = \"../tests/inputs/mobilenet.qat.onnx.engine\"\n",
    "plan = trex.EnginePlan(f\"{engine_name}.graph.json\", f\"{engine_name}.profile.json\", f\"{engine_name}.metadata.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e3de9cb",
   "metadata": {},
   "source": [
    "## List `k` slowest layers\n",
    "\n",
    "List the k-slowest layers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c601495",
   "metadata": {},
   "outputs": [],
   "source": [
    "top3 = plan.df.nlargest(3, 'latency.pct_time')\n",
    "for i in range(len(top3)):\n",
    "    layer = top3.iloc[i]\n",
    "    print(\"%s: %s\" % (layer[\"Name\"], layer[\"type\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd16e032",
   "metadata": {},
   "source": [
    "Compute the latency of the top-3 slowest layers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93d25555",
   "metadata": {},
   "outputs": [],
   "source": [
    "top3_latency = top3['latency.avg_time'].sum()\n",
    "top3_percent = top3['latency.pct_time'].sum()\n",
    "print(f\"top3 latency: {top3_latency:.6f} ms ({top3_percent:.2f}%)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "206b0f1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ltype = \"Convolution\"\n",
    "convs = plan.df.query(f\"type == \\\"{ltype}\\\"\")\n",
    "print(f\"There are {len(convs)} convolutions\")\n",
    "print(convs['latency.avg_time'].median())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a9c91f",
   "metadata": {},
   "outputs": [],
   "source": [
    "convs2 = plan.get_layers_by_type('Convolution')\n",
    "print(f\"There are {len(convs2)} convolutions\")\n",
    "print(convs['latency.avg_time'].median())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef4287a3",
   "metadata": {},
   "source": [
    "## Access layer activations\n",
    "\n",
    "There are several ways to access a layer's inputs and outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6cff068",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(convs.iloc[0]['Inputs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "216be540",
   "metadata": {},
   "outputs": [],
   "source": [
    "clean_convs = trex.clean_df(convs2.copy(), inplace=True)\n",
    "clean_convs.iloc[0]['Inputs']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e30298f8",
   "metadata": {},
   "source": [
    "As an `Activation` instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b482df",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, outputs = trex.create_activations(convs.iloc[0])\n",
    "print(inputs[0].name)\n",
    "print(inputs[0].shape)\n",
    "print(inputs[0].precision)\n",
    "print(inputs[0].format)\n",
    "print(inputs[0].size_bytes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e73772b7",
   "metadata": {},
   "source": [
    "## Query and Grouping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fbff2a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Group by type, and perform a sum reduction on the latency\n",
    "plan.df.groupby([\"type\"])[[\"latency.avg_time\", \"latency.pct_time\"]].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89985b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# trex provides another way to do the same thing\n",
    "trex.group_sum_attr(plan.df,\"type\", \"latency.avg_time\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "930d109f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Another trex convenience wrapper: group by 'type' and count the number of members in each group\n",
    "trex.group_count(plan.df, \"type\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "494e9968",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the dataframe of all INT8 convolutions with \n",
    "convs_1x1_fp32 = convs_1x1[convs_1x1[\"Outputs\"].str.startswith(\"FP32\")]\n",
    "display_df(convs_1x1_fp32)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
